"""
蒙特卡洛树搜索（MCTS）算法核心实现
包含Selection、Expansion、Simulation、Backpropagation四个步骤
"""

import random
import time
import numpy as np
from mcts_node import MCTSNode
from config import (
    MCTS_SIMULATIONS, 
    EXPLORATION_CONSTANT, 
    MAX_SIMULATION_DEPTH,
    WIN_COUNT,
    OPENING_BOOK
)


class MCTS:
    """蒙特卡洛树搜索算法类"""
    
    def __init__(self, simulations=MCTS_SIMULATIONS):
        """
        初始化MCTS
        
        Args:
            simulations: 每步模拟次数（越大越强但越慢）
        """
        self.simulations = simulations
        self.exploration_constant = EXPLORATION_CONSTANT
    
    def search(self, board, player):
        """
        执行MCTS搜索，返回最佳落子位置
        
        Args:
            board: 当前棋盘状态（numpy数组或Board对象）
            player: 当前玩家（1或-1）
        
        Returns:
            tuple: 最佳落子位置 (row, col)
        """
        # 检查是否使用开局库
        opening_move = self.get_opening_move(board)
        if opening_move:
            print(f"[MCTS] 使用开局库: {opening_move}")
            return opening_move
        
        # 转换为numpy数组
        if not isinstance(board, np.ndarray):
            board = board.copy()
        
        # 创建根节点
        root = MCTSNode(board, player=player)
        
        print(f"[MCTS] 开始搜索，模拟次数: {self.simulations}")
        start_time = time.time()
        
        # 执行指定次数的模拟
        for i in range(self.simulations):
            node = root
            
            # 1. Selection（选择）：从根节点开始，选择最有潜力的子节点
            while not node.is_terminal() and node.is_fully_expanded():
                node = node.select_child()
            
            # 2. Expansion（扩展）：在选中的节点上扩展一个新的子节点
            if not node.is_terminal() and not node.is_fully_expanded():
                node = node.expand()
            
            # 3. Simulation（模拟）：从新节点开始随机模拟到游戏结束
            result = self.simulate(node)
            
            # 4. Backpropagation（回溯）：将模拟结果回传到路径上的所有节点
            self.backpropagate(node, result)
            
            # 显示进度
            if (i + 1) % 100 == 0:
                elapsed = time.time() - start_time
                print(f"[MCTS] 进度: {i+1}/{self.simulations} ({elapsed:.1f}秒)")
        
        # 选择访问次数最多的子节点（最稳健的选择）
        if not root.children:
            # 如果没有子节点，返回任意合法落子
            legal_moves = root.get_legal_moves()
            return legal_moves[0] if legal_moves else (7, 7)
        
        best_child = max(root.children, key=lambda child: child.visits)
        
        elapsed = time.time() - start_time
        print(f"[MCTS] 搜索完成，耗时: {elapsed:.2f}秒")
        print(f"[MCTS] 最佳落子: {best_child.move}, 访问次数: {best_child.visits}, 胜率: {best_child.get_win_rate():.2%}")
        
        return best_child.move
    
    def simulate(self, node):
        """
        从当前节点随机模拟到游戏结束
        
        Args:
            node: 当前MCTS节点
        
        Returns:
            float: 模拟结果（1表示胜利，0表示平局，-1表示失败）
        """
        board = node.board.copy()
        current_player = node.player
        depth = 0
        
        while depth < MAX_SIMULATION_DEPTH:
            # 检查游戏是否结束
            winner = self.check_winner(board)
            if winner != 0:
                # 返回相对于根节点玩家的结果
                return 1.0 if winner == node.player else -1.0
            
            if self.is_board_full(board):
                return 0.0  # 平局
            
            # 随机选择一个合法落子
            legal_moves = self.get_legal_moves(board)
            if not legal_moves:
                return 0.0
            
            move = random.choice(legal_moves)
            board[move[0]][move[1]] = current_player
            
            # 切换玩家
            current_player = -current_player
            depth += 1
        
        # 达到最大深度，评估当前局势
        return self.evaluate_board(board, node.player)
    
    def backpropagate(self, node, result):
        """
        将模拟结果回传到路径上的所有节点
        
        Args:
            node: 叶子节点
            result: 模拟结果
        """
        while node is not None:
            node.update(result)
            result = -result  # 对手的结果相反
            node = node.parent
    
    def get_legal_moves(self, board):
        """
        获取合法落子位置（优化版：只考虑已有棋子周围）
        
        Args:
            board: 棋盘状态
        
        Returns:
            list: 合法落子位置列表
        """
        moves = []
        board_size = len(board)
        
        for i in range(board_size):
            for j in range(board_size):
                if board[i][j] == 0 and self.has_neighbor(board, i, j, distance=2):
                    moves.append((i, j))
        
        if not moves:
            # 如果没有找到，返回中心位置
            center = board_size // 2
            if board[center][center] == 0:
                moves.append((center, center))
        
        return moves
    
    def has_neighbor(self, board, row, col, distance=2):
        """
        检查周围是否有棋子
        
        Args:
            board: 棋盘状态
            row: 行坐标
            col: 列坐标
            distance: 搜索距离
        
        Returns:
            bool: 周围是否有棋子
        """
        board_size = len(board)
        for i in range(max(0, row - distance), min(board_size, row + distance + 1)):
            for j in range(max(0, col - distance), min(board_size, col + distance + 1)):
                if board[i][j] != 0:
                    return True
        return False
    
    def check_winner(self, board):
        """
        检查是否有玩家获胜
        
        Args:
            board: 棋盘状态
        
        Returns:
            int: 1（黑方胜）、-1（白方胜）、0（未结束）
        """
        board_size = len(board)
        directions = [(0, 1), (1, 0), (1, 1), (1, -1)]
        
        for i in range(board_size):
            for j in range(board_size):
                if board[i][j] == 0:
                    continue
                
                player = board[i][j]
                
                for dx, dy in directions:
                    count = 1
                    
                    # 正方向
                    x, y = i + dx, j + dy
                    while (0 <= x < board_size and 0 <= y < board_size 
                           and board[x][y] == player):
                        count += 1
                        x += dx
                        y += dy
                    
                    # 反方向
                    x, y = i - dx, j - dy
                    while (0 <= x < board_size and 0 <= y < board_size 
                           and board[x][y] == player):
                        count += 1
                        x -= dx
                        y -= dy
                    
                    if count >= WIN_COUNT:
                        return player
        
        return 0
    
    def is_board_full(self, board):
        """
        检查棋盘是否已满
        
        Args:
            board: 棋盘状态
        
        Returns:
            bool: 棋盘是否已满
        """
        return not np.any(board == 0)
    
    def evaluate_board(self, board, player):
        """
        简单的局势评估函数（用于达到最大深度时）
        
        Args:
            board: 棋盘状态
            player: 当前玩家
        
        Returns:
            float: 评估分数（-1到1之间）
        """
        # 简单实现：返回0（平局）
        # 可以实现更复杂的评估逻辑，如：
        # - 统计连子数量
        # - 评估威胁程度
        # - 计算位置价值
        return 0.0
    
    def get_opening_move(self, board):
        """
        使用开局库（前几步使用固定策略）
        
        Args:
            board: 棋盘状态
        
        Returns:
            tuple or None: 开局落子位置，如果不使用开局库则返回None
        """
        if not isinstance(board, np.ndarray):
            board = board.copy()
        
        move_count = np.sum(board != 0)
        
        # 第一步：天元
        if move_count == 0:
            return OPENING_BOOK['center']
        
        # 第二步：星位（随机选择）
        if move_count == 1:
            available_stars = [pos for pos in OPENING_BOOK['star_points'] 
                             if board[pos[0]][pos[1]] == 0]
            if available_stars:
                return random.choice(available_stars)
        
        return None
